import asyncio
import websockets
import json
import threading
import time
from time_util import datetime_str

clients = set()
lock = threading.Lock()
HEARTBEAT_INTERVAL = 10
main_loop = None
pending_responses = {}

# ---------------- 核心 WebSocket 逻辑 ----------------
async def send_heartbeat(websocket):
    while True:
        try:
            await websocket.send(json.dumps({"type": "heartbeat", "data": "pong"}))
            await asyncio.sleep(HEARTBEAT_INTERVAL)
        except Exception:
            break

async def client_handler(websocket, addr, window):
    log_message = f"[连接] 客户端 {addr}, {len(clients)} 已连接"
    print(f"{datetime_str()} - {log_message}")
    window.log_signal.emit(log_message)  # 发出信号，将日志信息传递到 UI
    with lock:
        clients.add(websocket)

    asyncio.create_task(send_heartbeat(websocket))

    try:
        async for message in websocket:
            try:
                data = json.loads(message)

                msg_type = data.get("type")
                payload = data.get("data")
                log_message = f"[消息] 客户端 {addr} 发送消息：{data}"
                print(f"{datetime_str()} - {log_message}")
                if msg_type != "heartbeat_ack" and msg_type!= "heartbeat":
                    
                    window.log_signal.emit(log_message)  # 发出信号，将日志信息传递到 UI

                if msg_type != "heartbeat_ack":
                    log_message = f"[消息] 客户端 {addr} 发送消息：{data}"
                    print(f"{datetime_str()} - {log_message}")
                    window.log_signal.emit(log_message)  # 发出信号，将日志信息传递到 UI
                if msg_type == "connect":
                    pass
                elif msg_type == "heartbeat"  or msg_type == "heartbeat_ack":  
                    pass
                elif msg_type == "task_ok":
                    req_id = data.get("id")
                    if req_id in pending_responses:
                        pending_responses[req_id].set_result(msg_type)
                elif msg_type == "task_error":
                    req_id = data.get("id")
                    if req_id in pending_responses:
                        pending_responses[req_id].set_result(msg_type)
                elif msg_type == "ack_point":
                    req_id = data.get("id")
                    if req_id in pending_responses:
                        pending_responses[req_id].set_result(payload)
                elif msg_type == "ack_field_text":
                    req_id = data.get("id")
                    if req_id in pending_responses:
                        pending_responses[req_id].set_result(payload)
                elif msg_type == "ping":
                    await websocket.send(json.dumps({"type": "pong"}))
                elif msg_type == "echo":
                    await websocket.send(json.dumps({"type": "echo", "data": payload}))
                elif msg_type == "broadcast":
                    await broadcast(f"[广播] {payload}")
                elif msg_type == "close":
                    await websocket.close()
                else:
                    await websocket.send(json.dumps({"type": "error", "data": "未知消息类型"}))

            except json.JSONDecodeError:
                await websocket.send(json.dumps({"type": "error", "data": "消息格式错误"}))

    except websockets.exceptions.ConnectionClosed:
        pass
    finally:
        with lock:
            clients.discard(websocket)
        log_message = f"[断开] 客户端 {addr} 已断开"
        print(f"{datetime_str()} - {log_message}")
        window.log_signal.emit(log_message)  # 发出信号，将日志信息传递到 UI

async def broadcast(message):
    with lock:
        active_clients = list(clients)
    if not active_clients:
        return
    await asyncio.gather(*[
        client.send(json.dumps({"type": "broadcast", "data": message}))
        for client in active_clients
    ])

async def broadcast_with_response(message, expect_type=None, timeout=5):
    request_id = str(time.time())
    payload = {
        "type": "command",
        "id": request_id,
        "data": message
    }
    print(f"{datetime_str()} - 发送消息：{payload}")
    with lock:
        active_clients = list(clients)
    if not active_clients:
        print(f"[WS_ERROR] 没有可用的客户端")
        return None

    future = asyncio.get_event_loop().create_future()
    pending_responses[request_id] = future
    print(f"{datetime_str()} - [WS_INFO] 等待响应：{request_id}")
    await asyncio.gather(*[
        client.send(json.dumps(payload))
        for client in active_clients
    ])

    try:
        result = await asyncio.wait_for(future, timeout=timeout)
        print(f"{datetime_str()} - [WS_INFO] 收到响应：{result}")
        return result
    except asyncio.TimeoutError:
        print(f"{datetime_str()} - [WS_ERROR] 等待响应超时：{request_id}")
        return None
    finally:
        pending_responses.pop(request_id, None)

async def connection_handler(websocket, window):
    addr = websocket.remote_address
    await client_handler(websocket, addr, window)

async def ws_main(window):
    """主 WebSocket 协程入口"""
    global main_loop
    main_loop = asyncio.get_running_loop()
    server = await websockets.serve(lambda ws: connection_handler(ws, window), "0.0.0.0", 8765)
    log_message = f"✅ WebSocket 服务已启动：ws://0.0.0.0:8765"
    print(f"{datetime_str()} - {log_message}")
    window.log_signal.emit(log_message)  # 发出信号，将日志信息传递到 UI
    await server.wait_closed()

# ---------------- 对外接口 ----------------
def has_clients():
    return bool(clients)

def send_to_all_clients(message: str):
    if main_loop is not None:
        future = asyncio.run_coroutine_threadsafe(broadcast(message), main_loop)
        return future.result(timeout=3)
    else:
        print(f"{datetime_str()} - ❌ 事件循环未初始化")

def send_and_wait_response(message: str, timeout=5):
    if main_loop is not None:
        future = asyncio.run_coroutine_threadsafe(
            broadcast_with_response(message, timeout=timeout),
            main_loop
        )
        return future.result(timeout=timeout + 1)
    else:
        print(f"{datetime_str()} - ❌ 事件循环未初始化")
        return None

def start_server_in_thread(window = None):
    """在独立线程中启动 WebSocket 服务器（安全，不会影响 PyQt）"""
    def _run():
        try:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            loop.run_until_complete(ws_main(window))
        except Exception as e:
            print(f"{datetime_str()} - [WS_THREAD_ERROR] {e}")
        finally:
            loop.close()

    t = threading.Thread(target=_run, daemon=True)
    t.start()
    return t

# ---------------- 测试入口 ----------------
if __name__ == "__main__":
    # 独立运行模式
    start_server_in_thread()
    while True:
        time.sleep(5)
        send_to_all_clients(f"系统通知：{time.strftime('%H:%M:%S')}")
